﻿#include	"Engine.TCP.hpp"
#include	"Engine.Crypto.hpp"
#include	"Engine.String.hpp"
#include	"Engine.Log.hpp"

#include	<cstring>

namespace Engine { namespace TCP {

#pragma		pack(push, 1)

	/// @Brief	: 发数据的协议头，每个协议之前会增加该结构体，发送时自动填充
	struct ProtoHeader {
		char		pMD5[16];	//!	整个除去pMD5的整个数据内容加上密文的签名，防止数据被修改
		uint64_t	nSerial;	//! 发送内容流水号，每个客户端从0开始只增，防止网络抓包重发
		uint32_t	nPF;		//! 记录一下协议簇，方便支持ProtocolBuffer等序列化库
		uint32_t	nSize;		//!	内容总大小（包括Header），用于网络分包
	};

	/// @Brief	: 协议信息。用于缓存收到协议
	struct ProtoStore {
		uint64_t	nConnId;	//!	记录消息来源，客户端缓存服务器消息时，该字段无效
		uint32_t	nPF;		//!	记录一下协议簇，方便支持ProtocolBuffer等序列化库
		uint32_t	nSize;		//!	内容总大小（包括Header），用于网络分包
	};

#pragma		pack(pop)

	//////////////////////////////////////////////////////////////////////////////
	/// Engine::TCP::IClient
	//////////////////////////////////////////////////////////////////////////////

	IClient::IClient() : _pClient(Network::Socket::Create()), _iCache(8 * 1024 * 1024), _nSerial(0) {
		::memset(_pToken, 0, 16);
		::memset(_szLogger, 0, 64);
	}

	IClient::~IClient() {
		Close();
	}

	void IClient::OnConnect(uint32_t uRemoteIp, uint16_t uPort) {
		std::string szIp = Network::GetIpString(uRemoteIp);
		int nPort = uPort;
		GLogger(_szLogger).LogInfo("Connected to server %s:%d successfully", szIp.c_str(), nPort);
	}

	bool IClient::OnReceive(char * pData, size_t uSize) {
		static unsigned char pMD5[16] = { 0 };
		static char pBuf[1024 * 256] = { 0 };

		size_t nRead = 0;
		while (nRead < uSize) {
			if (uSize - nRead <= sizeof(ProtoHeader)) {
				int nSize = uSize - nRead;
				GLogger(_szLogger).LogError("Received invalid protocol (size <= ProtoHeader). Size %d", nSize);
				break;
			}

			ProtoHeader * pHeader = (ProtoHeader *)(pData + nRead);
			nRead += pHeader->nSize + sizeof(ProtoHeader);

			::memcpy(pMD5, pHeader->pMD5, 16);
			::memcpy(pHeader->pMD5, _pToken, 16);
			
			Crypto::MD5 iMd5 = Crypto::CalcMD5(pData, pHeader->nSize + sizeof(ProtoHeader));
			if (::memcmp(iMd5.pData, pMD5, 16) != 0) {
				GLogger(_szLogger).LogError("Received invalid protocol (MD5 check). PF %d. MD5 %s", (int)pHeader->nPF, iMd5.ToString().c_str());
				continue;
			}

			::memset(pBuf, 0, 1024 * 256);
			::memcpy(pBuf + sizeof(ProtoStore), (char *)pHeader + sizeof(ProtoHeader), pHeader->nSize);

			ProtoStore * pStore = (ProtoStore *)pBuf;
			pStore->nConnId	= 0;
			pStore->nPF		= pHeader->nPF;
			pStore->nSize	= pHeader->nSize;

			if (!_iCache.Write(pStore, pHeader->nSize + sizeof(ProtoStore))) {
				GLogger(_szLogger).LogError("Alloc memory error! PF : %d, Size : %d", (int)pHeader->nPF, (int)pHeader->nSize);
				continue;
			}
		}

		return true;
	}

	void IClient::OnClose(int nReason) {
		GLogger(_szLogger).LogInfo("Disconnect from server. Reason : %d", nReason);
	}

	void IClient::Init(const char * pToken, const char * szLogger) {
		Crypto::MD5 iMd5 = Crypto::CalcMD5(pToken, strlen(pToken));
		::memcpy(_pToken, iMd5.pData, 16);
		String::Copy(_szLogger, szLogger, 64);
	}

	int IClient::Connect(const char * szIP, int nPort) {
		return _pClient->Connect(szIP, (uint16_t)nPort, this);
	}

	void IClient::Breath() {
		char *		pData	= _iCache.Read();
		uint32_t *	pSize	= (uint32_t *)pData;
		uint32_t	nRead	= 0;

		pData += sizeof(uint32_t);

		while (nRead < *pSize) {
			ProtoStore * pStore = (ProtoStore *)(pData + nRead);
			ProcessMessage(pStore->nPF, pData + nRead + sizeof(ProtoStore), pStore->nSize);
			nRead += pStore->nSize + sizeof(ProtoStore);
		}
	}

	int IClient::Send(int nPF, void * pData, uint32_t uSize) {
		static char pSendBuf[1024 * 256] = { 0 };

		if (uSize > 1024 * 256) {
			GLogger(_szLogger).LogError("Bad data size. Want to send bytes %d(>256KB) of data.", (int)uSize);
			return Network::ErrAlloc;
		}

		::memset(pSendBuf, 0, 1024 * 256);
		::memcpy(pSendBuf + sizeof(ProtoHeader), pData, uSize);

		ProtoHeader * pMsg = (ProtoHeader *)pSendBuf;

		pMsg->nSerial	= (++_nSerial);
		pMsg->nPF		= (uint32_t)nPF;
		pMsg->nSize		= uSize;

		::memcpy(pMsg->pMD5, _pToken, 16);
		Crypto::MD5 iMd5 = Crypto::CalcMD5((const char *)pMsg, uSize + sizeof(ProtoHeader));
		::memcpy(pMsg->pMD5, iMd5.pData, 16);

		return _pClient->Send(pSendBuf, uSize + sizeof(ProtoHeader));
	}

	void IClient::Close() {
		if (_pClient) Network::Socket::Close(_pClient, Network::ClientClose);
		_pClient = NULL;
	}

	//////////////////////////////////////////////////////////////////////////////
	/// Engine::TCP::IServer
	//////////////////////////////////////////////////////////////////////////////
	IServer::IServer()
		: _mpConns()
		, _pServer(Network::ServerSocket::Create())
		, _iCache(1024 * 1024 * 8) {
		::memset(_pToken, 0, 16);
		::memset(_szLogger, 0, 64);
	}

	IServer::~IServer() {
		Shutdown();
	}

	void IServer::OnAccept(uint64_t nConnId, uint32_t uRemoteIp, uint16_t uPort, int nState) {
		std::string szIP = Network::GetIpString(uRemoteIp);
		int nPort = uPort;

		Connection iConn;
		iConn.nConnId	= nConnId;
		iConn.nRemoteIp	= uRemoteIp;
		iConn.nPort		= uPort;
		iConn.nSerial	= 0;
		_mpConns.insert(std::pair<uint64_t, Connection>(nConnId, iConn));

		GLogger(_szLogger).LogInfo("Receive connect request from (%s:%d). ConnId %llu, State %d", szIP.c_str(), nPort, nConnId, nState);		
	}

	bool IServer::OnReceive(uint64_t nConnId, char * pData, size_t uSize) {
		static unsigned char pMD5[16] = { 0 };
		static char pBuf[1024 * 256] = { 0 };

		size_t nRead = 0;

		auto it = _mpConns.find(nConnId);
		if (it == _mpConns.end()) {
			GLogger(_szLogger).LogError("Receive invalid protocol from client %llu. Unknow client", nConnId);
			return false;
		}

		Connection iConn = it->second;

		while (nRead < uSize) {
			if (uSize - nRead <= sizeof(ProtoHeader)) {
				int nSize = uSize - nRead;
				GLogger(_szLogger).LogError("Received invalid protocol from client %llu (size <= ProtoHeader). Size %d", nConnId, nSize);
				break;
			}

			ProtoHeader * pHeader = (ProtoHeader *)(pData + nRead);
			nRead += pHeader->nSize + sizeof(ProtoHeader);

			::memcpy(pMD5, pHeader->pMD5, 16);
			::memcpy(pHeader->pMD5, _pToken, 16);

			Crypto::MD5 iMd5 = Crypto::CalcMD5(pData, pHeader->nSize + sizeof(ProtoHeader));
			if (::memcmp(pMD5, iMd5.pData, 16) != 0) {
				GLogger(_szLogger).LogError("Received invalid protocol from client %llu (MD5 check). PF %d. MD5 %s", nConnId, (int)pHeader->nPF, iMd5.ToString().c_str());
				continue;
			}

			if (iConn.nSerial >= pHeader->nSerial) {
				GLogger(_szLogger).LogError("Received invalid protocol from client %llu (Serial number). PF %d. Serial (%llu <=> %llu)", nConnId, (int)pHeader->nPF, pHeader->nSerial, iConn.nSerial);
				continue;
			}

			iConn.nSerial = pHeader->nSerial;
			
			::memset(pBuf, 0, 1024 * 256);
			::memcpy(pBuf + sizeof(ProtoStore), (char *)pHeader + sizeof(ProtoHeader), pHeader->nSize);

			ProtoStore * pStore = (ProtoStore *)pBuf;
			pStore->nConnId	= nConnId;
			pStore->nPF		= pHeader->nPF;
			pStore->nSize	= pHeader->nSize;

			if (!_iCache.Write(pBuf, pHeader->nSize + sizeof(ProtoStore))) {
				GLogger(_szLogger).LogError("Alloc memory error for client %llu! PF : %d, Size : %d", nConnId, (int)pHeader->nPF, (int)pHeader->nSize);
				continue;
			}
		}

		return true;
	}

	void IServer::OnClose(uint64_t nConnId, int nReason) {
		GLogger(_szLogger).LogInfo("Disconnect from client %llu. Reason %d", nConnId, nReason);
		_mpConns.erase(nConnId);
	}

	void IServer::Init(const char * pToken, const char * szLogger) {
		Crypto::MD5 iMd5 = Crypto::CalcMD5(pToken, strlen(pToken));
		::memcpy(_pToken, iMd5.pData, 16);
		String::Copy(_szLogger, szLogger, 64);
	}

	int IServer::Listen(const char * szIP, int nPort) {
		return _pServer->Listen(szIP, (uint16_t)nPort, this);
	}

	void IServer::Close(uint64_t nConnId) {
		_pServer->Close(nConnId, Network::SystemClose);
	}

	void IServer::Shutdown() {
		if (_pServer) Network::ServerSocket::Shutdown(_pServer);
		_pServer = NULL;
	}

	void IServer::Breath() {
		char *		pData	= _iCache.Read();
		uint32_t *	pSize	= (uint32_t *)pData;
		uint32_t	nRead	= 0;

		pData += sizeof(uint32_t);

		while (nRead < *pSize) {
			ProtoStore * pStore = (ProtoStore *)(pData + nRead);
			ProcessMessage(pStore->nConnId, pStore->nPF, pData + nRead + sizeof(ProtoStore), pStore->nSize);
			nRead += pStore->nSize + sizeof(ProtoStore);
		}
	}

	int IServer::Send(uint64_t nConnId, int nPF, void * pData, uint32_t uSize) {
		static char pSendBuf[1024 * 256] = { 0 };

		::memset(pSendBuf, 0, 1024 * 256);
		::memcpy(pSendBuf + sizeof(ProtoHeader), pData, uSize);

		ProtoHeader * pMsg = (ProtoHeader *)pSendBuf;
		pMsg->nPF		= (uint32_t)nPF;
		pMsg->nSerial	= 0;
		pMsg->nSize		= uSize;

		::memcpy(pMsg->pMD5, _pToken, 16);
		Crypto::MD5 iMd5 = Crypto::CalcMD5((const char *)pMsg, uSize + sizeof(ProtoHeader));
		::memcpy(pMsg->pMD5, iMd5.pData, 32);

		return _pServer->Send(nConnId, pSendBuf, uSize + sizeof(ProtoHeader));
	}

	int IServer::Broadcast(int nPF, void * pData, uint32_t uSize) {
		static char pSendBuf[1024 * 256] = { 0 };

		::memset(pSendBuf, 0, 1024 * 256);
		::memcpy(pSendBuf + sizeof(ProtoHeader), pData, uSize);

		ProtoHeader * pMsg = (ProtoHeader *)pSendBuf;
		pMsg->nPF		= (uint32_t)nPF;
		pMsg->nSerial	= 0;
		pMsg->nSize		= uSize;

		::memcpy(pMsg->pMD5, _pToken, 16);
		Crypto::MD5 iMd5 = Crypto::CalcMD5((const char *)pMsg, uSize + sizeof(ProtoHeader));
		::memcpy(pMsg->pMD5, iMd5.pData, 32);

		return _pServer->Broadcast(pSendBuf, uSize + sizeof(ProtoHeader));
	}

	IServer::Connection * IServer::GetClient(uint64_t nConnId) {
		auto it = _mpConns.find(nConnId);
		if (it == _mpConns.end()) return nullptr;
		return &(it->second);
	}
}}
